{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f49c3b2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "# For licensing see accompanying LICENSE file.\n",
    "# Copyright (c) 2025 Apple Inc. Licensed under MIT License.\n",
    "#\n",
    "\n",
    "# This is an example code for protein structure prediction with SimpleFold"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "231f5c19",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import numpy as np\n",
    "from math import pow\n",
    "import py3Dmol\n",
    "from pathlib import Path\n",
    "from io import StringIO\n",
    "from Bio.PDB import PDBIO\n",
    "from Bio.PDB import MMCIFParser, Superimposer\n",
    "sys.path.append(str(Path(\"./src/simplefold\").resolve()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cd07672",
   "metadata": {},
   "outputs": [],
   "source": [
    "# following are example amino acid sequences:\n",
    "example_sequences = {\n",
    "    \"7ftv_A\": \"GASKLRAVLEKLKLSRDDISTAAGMVKGVVDHLLLRLKCDSAFRGVGLLNTGSYYEHVKISAPNEFDVMFKLEVPRIQLEEYSNTRAYYFVKFKRNPKENPLSQFLEGEILSASKMLSKFRKIIKEEINDDTDVIMKRKRGGSPAVTLLISEKISVDITLALESKSSWPASTQEGLRIQNWLSAKVRKQLRLKPFYLVPKHAEETWRLSFSHIEKEILNNHGKSKTCCENKEEKCCRKDCLKLMKYLLEQLKERFKDKKHLDKFSSYHVKTAFFHVCTQNPQDSQWDRKDLGLCFDNCVTYFLQCLRTEKLENYFIPEFNLFSSNLIDKRSKEFLTKQIEYERNNEFPVFD\",\n",
    "    \"8cny_A\": \"MGPSLDFALSLLRRNIRQVQTDQGHFTMLGVRDRLAVLPRHSQPGKTIWVEHKLINILDAVELVDEQGVNLELTLVTLDTNEKFRDITKFIPENISAASDATLVINTEHMPSMFVPVGDVVQYGFLNLSGKPTHRTMMYNFPTKAGQCGGVVTSVGKVIGIHIGGNGRQGFCAGLKRSYFAS\",\n",
    "    \"8g8r_A\": \"GTVNWSVEDIVKGINSNNLESQLQATQAARKLLSREKQPPIDNIIRAGLIPKFVSFLGKTDCSPIQFESAWALTNIASGTSEQTKAVVDGGAIPAFISLLASPHAHISEQAVWALGNIAGDGSAFRDLVIKHGAIDPLLALLAVPDLSTLACGYLRNLTWTLSNLCRNKNPAPPLDAVEQILPTLVRLLHHNDPEVLADSCWAISYLTDGPNERIEMVVKKGVVPQLVKLLGATELPIVTPALRAIGNIVTGTDEQTQKVIDAGALAVFPSLLTNPKTNIQKEATWTMSNITAGRQDQIQQVVNHGLVPFLVGVLSKADFKTQKEAAWAITNYTSGGTVEQIVYLVHCGIIEPLMNLLSAKDTKIIQVILDAISNIFQAAEKLGETEKLSIMIEECGGLDKIEALQRHENESVYKASLNLIEKYFS\",\n",
    "    \"8i85_A\": \"MGILQANRVLLSRLLPGVEPEGLTVRHGQFHQVVIASDRVVCLPRTAAAAARLPRRAAVMRVLAGLDLGCRTPRPLCEGSLPFLVLSRVPGAPLEADALEDSKVAEVVAAQYVTLLSGLASAGADEKVRAALPAPQGRWRQFAADVRAELFPLMSDGGCRQAERELAALDSLPDITEAVVHGNLGAENVLWVRDDGLPRLSGVIDWDEVSIGDPAEDLAAIGAGYGKDFLDQVLTLGGWSDRRMATRIATIRATFALQQALSACRDGDEEELADGLTGYR\",\n",
    "}\n",
    "seq_id = \"7ftv_A\"  # choose from example_sequences\n",
    "aa_sequence = example_sequences[seq_id]\n",
    "print(f\"Predicting structure for {seq_id} with {len(aa_sequence)} amino acids.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb3c92f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "simplefold_model = \"simplefold_3B\" # choose from 100M, 360M, 700M, 1.1B, 1.6B, 3B\n",
    "backend = \"mlx\" # choose from [\"mlx\", \"torch\"]\n",
    "ckpt_dir = \"artifacts\"\n",
    "output_dir = \"artifacts\"\n",
    "prediction_dir = f\"predictions_{simplefold_model}_{backend}\"\n",
    "output_name = f\"{seq_id}\"\n",
    "num_steps = 500 # number of inference steps for flow-matching\n",
    "tau = 0.05 # stochasticity scale\n",
    "plddt = True # whether to use pLDDT confidence module\n",
    "nsample_per_protein = 1 # number of samples per protein"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "928b5156",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set random seed for reproducibility\n",
    "import lightning.pytorch as pl\n",
    "pl.seed_everything(42, workers=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf944c60",
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.simplefold.wrapper import ModelWrapper, InferenceWrapper\n",
    "\n",
    "# initialize the folding model and pLDDT model\n",
    "model_wrapper = ModelWrapper(\n",
    "    simplefold_model=simplefold_model,\n",
    "    ckpt_dir=ckpt_dir,\n",
    "    plddt=plddt,\n",
    "    backend=backend,\n",
    ")\n",
    "device = model_wrapper.device\n",
    "folding_model = model_wrapper.from_pretrained_folding_model()\n",
    "plddt_model = model_wrapper.from_pretrained_plddt_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4e924c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize the inference module with inference configurations\n",
    "inference_wrapper = InferenceWrapper(\n",
    "    output_dir=output_dir,\n",
    "    prediction_dir=prediction_dir,\n",
    "    num_steps=num_steps,\n",
    "    tau=tau,\n",
    "    nsample_per_protein=nsample_per_protein,\n",
    "    device=device,\n",
    "    backend=backend\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaaae97d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# process input sequence and run inference\n",
    "batch, structure, record = inference_wrapper.process_input(aa_sequence)\n",
    "\n",
    "results = inference_wrapper.run_inference(\n",
    "    batch,\n",
    "    folding_model,\n",
    "    plddt_model,\n",
    "    device=device,\n",
    ")\n",
    "save_paths = inference_wrapper.save_result(\n",
    "    structure,\n",
    "    record,\n",
    "    results,\n",
    "    out_name=output_name\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bf2e7e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize the first predicted structure\n",
    "pdb_path = save_paths[0]\n",
    "view = py3Dmol.view(query=pdb_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6173d76d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# color based on the predicted confidence\n",
    "# confidence coloring from low to high: red–orange–yellow–green–blue (0 to 100)\n",
    "if plddt:\n",
    "    view.setStyle({'cartoon':{'colorscheme':{'prop':'b','gradient':'roygb','min':0,'max':100}}})\n",
    "    view.zoomTo()\n",
    "    view.show()\n",
    "# color in spectrum if pLDDT is not available\n",
    "else:\n",
    "    view.setStyle({'cartoon':{'color':'spectrum'}})\n",
    "    view.zoomTo()\n",
    "    view.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2bccf93",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize the all-atom structure\n",
    "view.setStyle({'stick':{}})\n",
    "view.zoomTo()\n",
    "view.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d918192",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize the predicted structure in 3D alongside the GT structure\n",
    "\n",
    "def calculate_tm_score(coords1, coords2, L_target=None):\n",
    "    \"\"\"\n",
    "    Compute TM-score for two aligned coordinate sets (numpy arrays).\n",
    "    \n",
    "    coords1, coords2: Nx3 numpy arrays (aligned atomic coordinates, e.g. CA atoms)\n",
    "    L_target: length of target protein (default = len(coords1))\n",
    "    \"\"\"\n",
    "    assert coords1.shape == coords2.shape, \"Aligned coords must have same shape\"\n",
    "    N = coords1.shape[0]\n",
    "\n",
    "    if L_target is None:\n",
    "        L_target = N\n",
    "\n",
    "    # distances between aligned atoms\n",
    "    dists = np.linalg.norm(coords1 - coords2, axis=1)\n",
    "\n",
    "    # scaling factor d0\n",
    "    d0 = 1.24 * pow(L_target - 15, 1/3) - 1.8\n",
    "    if d0 < 0.5:\n",
    "        d0 = 0.5  # safeguard, as in TM-align\n",
    "\n",
    "    # TM-score\n",
    "    score = np.sum(1.0 / (1.0 + (dists/d0)**2)) / L_target\n",
    "    return score\n",
    "\n",
    "parser = MMCIFParser(QUIET=True)\n",
    "\n",
    "# Load two structures\n",
    "struct1 = parser.get_structure(\"ref\", f\"assets/{seq_id}.cif\")\n",
    "struct2 = parser.get_structure(\"prd\", pdb_path)\n",
    "\n",
    "# Select CA atoms for alignment\n",
    "atoms1 = [a for a in struct1.get_atoms() if a.get_id() == 'CA']\n",
    "atoms2 = [a for a in struct2.get_atoms() if a.get_id() == 'CA']\n",
    "print(len(atoms1), len(atoms2))\n",
    "\n",
    "# Superimpose\n",
    "sup = Superimposer()\n",
    "sup.set_atoms(atoms1, atoms2)\n",
    "sup.apply(struct2.get_atoms())\n",
    "\n",
    "# Calculate TM-score\n",
    "coords1 = np.array([a.coord for a in atoms1])\n",
    "coords2 = np.array([a.coord for a in atoms2])\n",
    "tm_score = calculate_tm_score(coords1, coords2)\n",
    "\n",
    "print(\"TM-score (0-1, higher is better): {:.3f}\".format(tm_score))\n",
    "print(\"RMSD (lower is better): {:.3f}\".format(sup.rms))\n",
    "\n",
    "# Save aligned structures to strings\n",
    "io = PDBIO()\n",
    "\n",
    "s1_buf, s2_buf = StringIO(), StringIO()\n",
    "io.set_structure(struct1); io.save(s1_buf)\n",
    "io.set_structure(struct2); io.save(s2_buf)\n",
    "\n",
    "# Visualize in py3Dmol\n",
    "view = py3Dmol.view(width=600, height=400)\n",
    "view.addModel(s1_buf.getvalue(),\"pdb\")\n",
    "view.addModel(s2_buf.getvalue(),\"pdb\")\n",
    "\n",
    "# Color reference protein blue, predicted structure red\n",
    "view.setStyle({'model': 0}, {'cartoon': {'color': 'blue'}})\n",
    "view.setStyle({'model': 1}, {'cartoon': {'color': 'red'}})\n",
    "\n",
    "# Add legend\n",
    "view.addLabel(\"Ground Truth\", {'position': {'x': 0, 'y': 0, 'z': 0}, 'backgroundColor': 'blue', 'fontColor': 'white', 'fontSize': 12})\n",
    "view.addLabel(\"Predicted\", {'position': {'x': 0, 'y': 4, 'z': 0}, 'backgroundColor': 'red', 'fontColor': 'white', 'fontSize': 12})\n",
    "\n",
    "view.zoomTo()\n",
    "view.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f8709b6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "simplefold_mlx",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
